﻿# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.stats import rayleigh
from concurrent.futures import ProcessPoolExecutor
from functools import partial
# 导入定态薛定谔方程类
from quantum_mechanics.stationary_schroedinger import StationarySchroedingerEquation, infinite_well_potential

def compute_rayleigh_pdf(scale, x):
    """计算指定scale下的瑞丽分布PDF值"""
    return rayleigh.cdf(x, scale=scale)

def main1():
    my_points = np.array([
    [0, 0, 0],
    [3, 0, 0],
    [1, 2, 0],
    [0, 1, 4],
    [4, 2, 0],
    [3, 1, 4],
    [1, 3, 4],
    [4, 3, 4]

])
    res =plot_3d_points(my_points, colors='purple', sizes=100, title="我的三维点")
    res[0].savefig("my_3d_points.png", dpi=300, bbox_inches='tight')
    

def main():
    # 创建3D图形
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')

    # 定义相同的x定义域（增加采样点以显示并行计算优势）
    x = np.linspace(0, 5, 1000)  # 增加到1000个点

    # 设置更多不同的scale参数
    scales = np.linspace(0.5, 2.5, 30)  # 增加到20个scale值

    # 创建部分应用函数，固定x参数
    partial_compute = partial(compute_rayleigh_pdf, x=x)

    # 使用进程池并行计算
    with ProcessPoolExecutor() as executor:
        results = list(executor.map(partial_compute, scales))

    # 绘制结果
    for i, scale in enumerate(scales):
        y = results[i]
        # 在3D空间中绘制曲线，使用不同的高度表示不同的scale
        ax.plot(x, [scale] * len(x), y, alpha=0.6)

    # 设置标签
    ax.set_xlabel('X轴')
    ax.set_ylabel('Scale参数')
    ax.set_zlabel('Rayleigh CDF值')
    ax.set_title('不同Scale参数下的Rayleigh分布CDF曲线')

    # 添加网格线
    ax.grid(True)

    return fig, ax

def plot_3d_points(points, colors=None, sizes=None, title=None, save_path=None):
    """
    绘制三维点图
    
    参数:
    - points: 三维点的数组，形状为(N, 3)
    - colors: 点的颜色，字符串或数组
    - sizes: 点的大小，标量或数组
    - title: 图表标题
    - save_path: 保存路径，如果为None则不保存
    
    返回:
    - (fig, ax): 图形和坐标轴对象
    """
    # 创建图形
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 提取x, y, z坐标
    x = points[:, 0]
    y = points[:, 1]
    z = points[:, 2]
    
    # 绘制点
    scatter = ax.scatter(x, y, z, c=colors, s=sizes, alpha=0.7)
    
    # 设置标签
    ax.set_xlabel('X轴')
    ax.set_ylabel('Y轴')
    ax.set_zlabel('Z轴')
    
    # 设置标题
    if title is not None:
        ax.set_title(title)
    
    # 添加网格线
    ax.grid(True)


def plot_stationary_wavefunction():
    """
    定态波函数绘图
    
    功能: 给定x在区间0到1物理意义是1nm，将其分为1000份取点。
         计算出波函数的平方即概率作为三维z轴的值，时间t作为y轴的值。
         利用定态波函数中的类中的完整的波函数计算z值，绘制三维点呈现出来。
    """
    # 创建定态薛定谔方程实例
    se = StationarySchroedingerEquation()
    
    # 设置势能函数（使用无限深势阱）
    # 注意：由于x范围是0到1，我们需要调整无限深势阱的位置
    se.set_potential(lambda x: infinite_well_potential(x - 0.5, width=1.0, height=100.0))
    
    # 设置空间网格（0到1nm，分为1000份）
    se.set_space_grid(0, 1, 1000)
    
    # 求解前几个本征值和本征函数
    eigenvalues, eigenfunctions = se.solve_eigenvalues(n=1)
    
    # 获取第一个本征函数（基态）
    spatial_wavefunction = se.get_spatial_wavefunction(eigenfunctions[:, 0])
    energy = eigenvalues[0]
    
    # 定义时间范围和点数
    t_start, t_end = 0, 10  # 时间范围，单位可以是任意时间单位（已延长时间范围）
    t_points = 200  # 时间采样点数（相应增加以保持时间分辨率）
    t_values = np.linspace(t_start, t_end, t_points)
    
    # 获取空间位置数组
    x_start, x_end = se.x_range
    x_values = np.linspace(x_start, x_end, se.x_points)
    
    # 准备存储三维数据的数组
    points = []
    
    # 对每个时间点计算完整波函数并获取概率密度
    for t in t_values:
        # 构建完整的波函数（包含时间演化）
        full_wavefunction = se.construct_full_wavefunction(spatial_wavefunction, energy, t)
        
        # 计算波函数的平方作为概率密度
        probability_density = np.abs(full_wavefunction)**2
        
        # 将当前时间点的所有x和对应的概率密度添加到点集
        for i, x in enumerate(x_values):
            points.append([x, t, probability_density[i]])
    
    # 转换为numpy数组
    points_array = np.array(points)
    
    # 创建3D图形
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制点，使用颜色映射表示概率密度
    scatter = ax.scatter(
        points_array[:, 0], 
        points_array[:, 1], 
        points_array[:, 2], 
        c=points_array[:, 2], 
        cmap='viridis', 
        alpha=0.6, 
        s=2
    )
    
    # 添加颜色条
    cbar = plt.colorbar(scatter, ax=ax, shrink=0.7, aspect=10)
    cbar.set_label('概率密度 |Ψ(x,t)|²')
    
    # 设置标签和标题
    ax.set_xlabel('位置 x (nm)')
    ax.set_ylabel('时间 t')
    ax.set_zlabel('概率密度 |Ψ(x,t)|²')
    ax.set_title('定态波函数的时间演化与概率密度分布')
    
    # 设置视角
    ax.view_init(elev=20, azim=45)
    
    # 添加网格线
    ax.grid(True, alpha=0.3)
    
    # 显示图形
    plt.tight_layout()
    plt.show(block=True)  # 使用block=True确保图形窗口保持打开
    
    return fig, ax


# 如果直接运行此脚本，可以测试新添加的功能
if __name__ == "__main__":
    # 调用定态波函数绘图方法
    #plot_stationary_wavefunction()
    
    # 导入并使用一维无限深势阱类显示波函数、概率密度和能级图
    from quantum_mechanics.infinite_well import InfiniteSquareWell
    
    # 创建一维无限深势阱实例
    well = InfiniteSquareWell(length=1.0, mass=1.0, hbar=1.0)
    
    # 显示第1个能级的波函数
    well.plot_wavefunction(n=1)
    
    # 显示第1个能级的概率密度
    well.plot_probability_density(n=1)
    
    # 显示前5个能级的能量图
    well.plot_energy_levels(max_n=5)